#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd
from scipy import stats

GCM=sys.argv[1]  # Enter GCM
RCM=sys.argv[2]  # Enter RCM
rcp=sys.argv[3]  # Enter RCP scenario

path_ki='../0-DATA/K-index/'+RCM+'/'+GCM+'/'
path_outs='./'+RCM+'/'+GCM+'/'
path_thres='../1-FWI/'+RCM+'/'+GCM+'/'

from dask_mpi import initialize
initialize()

from distributed import Client
client = Client()

# Read files
kindexs=sorted(glob.glob(path_kindex+'Ki_historical_*.nc')+glob.glob(path_kindex+'Ki_'+rcp+'_*.nc'))
Ki = xr.open_mfdataset(kindexs,combine='by_coords')['Ki']
FWI_thres = xr.open_dataset(path_thres+'FWI_90threshold_'+rcp+'.nc')['fwi']
Ki_thres = xr.full_like(FWI_thres,25.)
Ki = Ki.sel(time=slice('1970-01-01','2098-12-31'))

num = xr.full_like(Ki, 0)           # Daily array full of 0
num = xr.where((Ki>Ki_thres),1,num) # Days with Ki above the threshold are classified as 1
ndays = num.chunk({'time':-1,'rlat': 15, 'rlon': 15}).resample(time='1Y').sum(dim='time')  # Number of days with Ki above the threshold
ndays.chunk({'time':-1,'rlat': 15, 'rlon': 15}).to_netcdf(path_outs+'Ki_ndays_'+rcp+'.nc') # Save output
